import os
from pprint import pprint

from zhipuai import ZhipuAI, APIRequestFailedError

client = ZhipuAI(api_key=os.getenv('ZHIPUAI_API_KEY'))  # 请填写您自己的APIKey

def zhipu_call(model="GLM-4-Air-250414", convs=None):
    if convs is None:
        convs = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "中国的首都是是哪？网上查一下，然后简短的告诉我。"}
        ]
    flag = 0
    try:
        response = client.chat.completions.create(
            model=model,
            messages=convs,
            tools=[{
                'type': 'web_search',
                'web_search': {
                    'enable': True,
                    'content_size': 'high',
                    'search_result': True,
                    'require_search': True,
                }
            }]
        )
    except APIRequestFailedError:
        flag = 1
        response = client.chat.completions.create(
            model=model,
            messages=convs,
        )
    if flag == 0:
        return_body = {
            "search": response.web_search,
            "thinking": response.choices[0].message.reasoning_content,
            "answer": response.choices[0].message.content,
            "usage": {
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens
            }
        }
    else:
        return_body = {
            "search": [],
            "thinking": response.choices[0].message.reasoning_content,
            "answer": response.choices[0].message.content,
            "usage": {
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens
            }
        }
    return return_body

if __name__ == '__main__':
    res = zhipu_call()
    pprint(res)
